#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/4/21 16:18
@Author  : HZP
@File    : 2.JsonOutputParser使用技巧.py
"""
from dotenv import load_dotenv
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

load_dotenv()


class Joke(BaseModel):
    joke: str = Field(description="回答用户的冷笑话")
    punchline: str = Field(description="这个冷笑话的笑点")


json_out = JsonOutputParser(pydantic_object=Joke)

prompt = ChatPromptTemplate.from_template(
    "请根据用户的提问进行回答。\n{format_instructions}\n{query}").partial(
    format_instructions=json_out.get_format_instructions())

llm = ChatOpenAI(
    temperature=0.6,
    model="deepseek-r1:70b"
)
ai_message = json_out.invoke(
    llm.invoke(prompt.invoke({"query": "请给我讲一个冷笑话"})))

print(ai_message.get("joke"))
print(ai_message.get("punchline"))
